import numpy as np
from pylab import *
from matplotlib import pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.image as mpimg
from Interpolation import *
from Sector_Parameters import *

font = {'family' : 'monospace', 'weight' : 'bold', 'size'   : 40}
mpl.rc('font', **font)  # pass in the font dict as kwargs
mpl.rcParams['axes.linewidth'] = 10

Nr,Nt,Nz,dr,dt,dz,Np,NN,Nx,Nw,Ny,Nv,N_p, Nx_aug,sigma_pnts=circular_parameters()
DeltaT, Tsim, Nsim=time_parameters()

Nt_new = 85
Nr_new = 146

def surface_parameters(d,par,str1, filename):
    P=[]
    for j in range(Nt):
        c=d[int(((Nt+j)-Nt)*Nr):int(((Nt+j+1)-Nt)*Nr)] 
        P.append(c)
    PP=np.array(P)
    z=PP
    z_new=new_values(z)
    z_new_ = new_values_radial(z_new)
    Z_new = z_new_.squeeze()
    Z_new = np.transpose(Z_new)
    fig = plt.figure()
    r_new=linspace(0,290,Nr_new)
    th_new = linspace(pi, pi/2, Nt_new)
    R, TH = meshgrid(r_new, th_new)
    m=plt.subplot(projection="polar")
    m.set_xticklabels([])
    xlim([pi,pi/2])# Creating a sector from the 
    m.set_yticklabels([])
    plt.pcolormesh(TH,R,Z_new,cmap='turbo',shading='gouraud')
    plt.colorbar(label=str1)
    plt.clim(Z_new.min(),Z_new.max())
    fig.tight_layout()
    plt.savefig(par+'_'+filename+'.pdf',dpi=300)
    
    return Z_new

def surface_parameters_orig(d,par,str1, filename):
    Z_new = d
    fig=plt.figure(figsize=(25,25))
    ax=Axes3D(fig)
    r_new=linspace(0, 290, Nr_new)
    th_new = np.linspace(np.pi, np.pi/2, Nt_new)
    R, TH = meshgrid(r_new, th_new)
    m=plt.subplot(projection="polar")
    m.set_xticklabels([])
    m.grid(linewidth=1.5, color='k')
    xlim([pi,pi/2])
    m.set_yticklabels([])
    colorList = ['#999900', '#33CC00', '#3366CC']
    cmap = matplotlib.colors.ListedColormap(colorList)
    bounds = [0.,1.,2.,3.]
    ticks = [0.5,1.5,2.5]
    
    p = plt.pcolormesh(TH,R,Z_new,cmap=cmap, vmin = 0, vmax=3, shading='flat')
    cbar = fig.colorbar(p, aspect = 20, ticks=ticks, shrink=0.8)
    cbar.set_ticklabels(['MZ$_1$','MZ$_2$','MZ$_3$'])
    cbar.set_label(label='Legend', weight='bold', family='monospace', size=40, labelpad= 14)
    fig.tight_layout()
    plt.savefig('./results/MZs_Q4.jpg')
    plt.show()